
#include "zrpc/rpc/rpc_channel.h"

#include <google/protobuf/descriptor.h>

#include "zrpc/base/logger.h"

const char rpctag[] = "RPC0";
using namespace zrpc;

RpcChannel::RpcChannel()
    : codec_(std::bind(&RpcChannel::OnRpcMessage, this, std::placeholders::_1,
                       std::placeholders::_2)),
      services_(nullptr) {}

RpcChannel::RpcChannel(const std::shared_ptr<TcpConnection>& conn)
    : codec_(std::bind(&RpcChannel::OnRpcMessage, this, std::placeholders::_1,
                       std::placeholders::_2)),
      conn_(conn),
      services_(nullptr) {}

RpcChannel::~RpcChannel() {
  for (const auto& outstanding : outstandings_) {
    OutstandingCall out = outstanding.second;
    delete out.response;
    delete out.done;
  }
}

void RpcChannel::CallMethod(const ::google::protobuf::MethodDescriptor* method,
                            google::protobuf::RpcController* controller,
                            const ::google::protobuf::Message* request,
                            ::google::protobuf::Message* response,
                            ::google::protobuf::Closure* done) {
  RpcMessage message;
  message.set_type(REQUEST);
  int64_t id = id_++;
  message.set_id(id);
  message.set_service(method->service()->full_name());
  message.set_method(method->name());
  message.set_request(request->SerializeAsString());

  OutstandingCall out = {response, done};
  {
    std::unique_lock<std::mutex> lk(mutex_);
    outstandings_[id] = out;
  }
  codec_.Send(conn_, message);
}

void RpcChannel::OnMessage(const std::shared_ptr<TcpConnection>& conn,
                           Buffer* buf) {
  codec_.OnMessage(conn, buf);
}

void RpcChannel::OnRpcMessage(const std::shared_ptr<TcpConnection>& conn,
                              const std::shared_ptr<RpcMessage>& message) {
  assert(conn == conn_);
  if (message->type() == RESPONSE) {
    int64_t id = message->id();
    OutstandingCall out = {nullptr, nullptr};

    {
      std::unique_lock<std::mutex> lk(mutex_);
      auto it = outstandings_.find(id);
      if (it != outstandings_.end()) {
        out = it->second;
        outstandings_.erase(it);
      }
    }

    if (out.response) {
      std::unique_ptr<google::protobuf::Message> d(out.response);
      out.response->ParseFromString(message->response());
      if (out.done) {
        out.done->Run();
      }
    }
  } else if (message->type() == REQUEST) {
    ErrorCode error = WRONG_PROTO;
    if (services_) {
	  const auto it =
          services_->find(message->service());
      if (it != services_->end()) {
        google::protobuf::Service* service = it->second;
        assert(service != nullptr);
        const google::protobuf::ServiceDescriptor* desc =
            service->GetDescriptor();
        const google::protobuf::MethodDescriptor* method =
            desc->FindMethodByName(message->method());
        if (method) {
          std::unique_ptr<google::protobuf::Message> request(
              service->GetRequestPrototype(method).New());
          if (request->ParseFromString(message->request())) {
            google::protobuf::Message* response =
                service->GetResponsePrototype(method).New();
            int64_t id = message->id();
            service->CallMethod(
                method, nullptr, request.get(), response,
                NewCallback(this, &RpcChannel::DoneCallback, response, id));
            error = NO_ERROR;
          } else {
            error = INVALID_REQUEST;
          }
        } else {
          error = NO_METHOD;
        }
      } else {
        error = NO_SERVICE;
      }
    } else {
      error = NO_SERVICE;
    }

    if (error != NO_ERROR) {
      RpcMessage response;
      response.set_type(RESPONSE);
      response.set_id(message->id());
      response.set_error(error);
      codec_.Send(conn_, response);
    }
  } else if (message->type() == ERROR) {
  }
}

void RpcChannel::DoneCallback(::google::protobuf::Message* response,
                              int64_t id) {
  std::unique_ptr<google::protobuf::Message> d(response);
  RpcMessage message;
  message.set_type(RESPONSE);
  message.set_id(id);
  message.set_response(response->SerializeAsString());
  codec_.Send(conn_, message);
}
